import os
import time
import numpy as np
import tensorflow as tf
from model import Encode
import modeling

pre_trained_model = '../../../CodeLearning/resource/pretrained-bert-model/chinese_L-12_H-768_A-12'

config = os.path.join(pre_trained_model, 'bert_config.json')
checkpoint_path = os.path.join(pre_trained_model, 'bert_model.ckpt')
vocab_file = os.path.join(pre_trained_model, 'vocab.txt')

bert_config = modeling.BertConfig.from_json_file(config)


testText = ['何凯是个G!']


if __name__ == '__main__':
    conf = tf.ConfigProto(
        allow_soft_placement=True,
        log_device_placement=False,
        gpu_options=tf.GPUOptions(allow_growth=True))
    with tf.Session(config=conf) as sess:
        model = Encode(bert_config, vocab_file)
        model.init(checkpoint_path)
        sess.run(tf.global_variables_initializer())

        for sent in testText:
            feature = model.encode(sess, sent)

            print(feature)
            np.save('feature1.npy', feature)
